#pragma once

#include <cstdint>
#include <cstring>
#include <vector>
#include <cassert>
#include <unordered_set>
#include <map>
#include "util.h"

static constexpr int INTIAL_FANOUT   = 256;
static constexpr int MAX_FANOUT      = 4096;
static constexpr int MAX_LEAF        = 32;
static constexpr int AVERAGE_KEYLEN  = 64;
static constexpr float RESIZE_FACTOR = 2;
static constexpr int INLINE_PREFIX   = 16;
static const auto Hash1 = util::BKDRHash;
static const auto Hash2 = util::APHash;

struct MyKey {
    int len;
    char * str;

    struct KeyEqn {
        bool operator()(const MyKey& a, const MyKey& b) const{
            return memcmp(a.str, b.str, a.len) == 0;
        }
    };

    struct KeyHash {
        int operator()(const MyKey &a) const {
            return Hash1(a.str, a.len);
        }
    };
};

struct KeyTable {
    int capacity_;
    int16_t count_; 
    int16_t span_;
    char * buffer_;

    KeyTable(uint16_t span) : span_(span) {
        capacity_ = INTIAL_FANOUT * span_;
        buffer_ = new char[capacity_];
        count_ = 0;
    }

    KeyTable(std::vector<char *> & init_keys, uint16_t span) : span_(span) {
        capacity_ = init_keys.size() * span_ * RESIZE_FACTOR;
        buffer_ = new char[capacity_];
        
        for(int i = 0; i < init_keys.size(); i++) {
            memcpy(buffer_ + i * span_, init_keys[i], span_);
        }
        count_ = init_keys.size();
    }

    ~KeyTable() {
        delete [] buffer_;
    }

    void Insert(const char * piece) {
        if(capacity_ == count_ * span_) {
            ResizeAndInsert(piece);
        } else {
            int i = 0;
            for(; i < count_; i++) {
                if(memcmp(piece, buffer_ + i * span_, span_) < 0) {
                    break;
                }
            }

            memmove(buffer_ + (i + 1) * span_, buffer_ + i * span_, (count_ - i) * span_);
            memcpy(buffer_ + i * span_, piece, span_);
            count_ += 1;
        }
    }


    void ResizeAndInsert(const char * piece) {
        int new_cap = capacity_ * RESIZE_FACTOR;
        char * new_buffer = new char[new_cap];
        int i = 0;
        for(; i < count_; i++) {
            if(memcmp(piece, buffer_ + i * span_, span_) < 0) {
                break;
            } else {
                memcpy(new_buffer + i * span_, buffer_ + i * span_, span_);
            }
        }
        memcpy(new_buffer + i * span_, piece, span_);
        memcpy(new_buffer + (i + 1) * span_, buffer_ + i * span_, (count_ - i) * span_);
        // update meta data
        std::swap(new_buffer, buffer_);
        std::swap(new_cap, capacity_);
        count_ += 1;
        delete [] new_buffer;
    }

    char * operator[](uint16_t idx) {
        assert(idx < count_);
        return buffer_ + idx * span_;
    }
};

struct HashTable {
    int capacity_;
    int16_t count_;
    int16_t span_;
    uint64_t *array_;
    
    HashTable(uint16_t span, int cap = INTIAL_FANOUT) : span_(span) {
        capacity_ = INTIAL_FANOUT;
        count_ = 0;
        array_ = new uint64_t[capacity_]; 
    }

    HashTable(std::vector<char *> & init_keys, uint16_t span, 
                std::vector<void*> & init_childs) : span_ (span) {
        capacity_ = init_keys.size() * 2;
        array_ = new uint64_t[capacity_];

        for(int i = 0; i < init_keys.size(); i++) {
            Insert(init_keys[i], init_childs[i]);
        }
    }

    ~HashTable () {
        delete array_;
    }

    // looking for mapping associated with key
    void * Lookup(const char * piece) const {
        uint16_t hash_id1 = Hash1(piece, span_) % capacity_;
        uint16_t hash_id2 = Hash2(piece, span_) % capacity_;
        if((array_[hash_id1] >> 48) == hash_id2) {
            uint64_t pointer = array_[hash_id1] & 0x0000ffffffffffff;
            return (void *) pointer;
        } else if((array_[hash_id2] >> 48) == hash_id1) {
            uint64_t pointer = array_[hash_id2] & 0x0000ffffffffffff;
            return (void *) pointer;
        } else {
            return nullptr;
        }
    }

    // insert the key => child mapping into array
    bool Insert(const char * piece, void * child) {
        uint16_t hash_id1 = Hash1(piece, span_) % capacity_;
        uint16_t hash_id2 = Hash2(piece, span_) % capacity_;
        if(array_[hash_id1] == 0) {
            uint64_t value = ((uint64_t) hash_id2 << 48) + (uint64_t) child;
            array_[hash_id1] = value;
            
            count_ += 1;
            return true;
        } else if(array_[hash_id2] == 0) {
            uint64_t value = ((uint64_t) hash_id1 << 48) + (uint64_t) child;
            array_[hash_id2] = value;

            count_ += 1;
            return true;
        } else {
            // now we first choose the second one as victim
            uint16_t repeat_times = 0;
            uint16_t old_value = ((uint64_t) hash_id1 << 48) + (uint64_t)child;
            uint16_t victim = hash_id2;
            uint16_t next_victim = array_[hash_id2] >> 48;
            while(repeat_times < 5) {
                uint64_t victim_value = ((uint64_t) victim << 48) + array_[victim] & 0x0000ffffffffffff;
                array_[victim] = old_value;
                if(array_[next_victim] == 0) {
                    array_[next_victim] = victim_value;

                    count_ += 1;
                    return true;
                } else {
                    old_value = victim_value;
                    victim = next_victim;
                    next_victim = array_[next_victim] >> 48;
                    repeat_times++;
                }
            } 
            return false;
        }
    }
};

class TrieNode {
private:
    uint16_t type;
    uint16_t prefix_len;
    uint16_t trie_span;
    uint16_t count;
    
    char inline_prefix[INLINE_PREFIX];
    char * prefix_suffix;

    // partial keys and corresponding children
    KeyTable keys;
    HashTable mappings;

public:
    TrieNode(uint16_t span, const char * p = nullptr) : keys(span), mappings(span) {
        prefix_len = strlen(p);
        trie_span = span;
        count = 0;

        memcpy(inline_prefix, p, (prefix_len <= INLINE_PREFIX ? prefix_len : INLINE_PREFIX));
        if(prefix_len > INLINE_PREFIX) {
            prefix_suffix = new char[prefix_len - INLINE_PREFIX];
            memcpy(prefix_suffix, p + INLINE_PREFIX, prefix_len - INLINE_PREFIX);
        } else {
            prefix_suffix = nullptr;
        }
    }

    ~TrieNode() {
        if (prefix_suffix) 
            delete [] prefix_suffix;
    }

    /// @brief Find a child with matching key
    /// @param piece the lookup key prefix
    /// @param piece_len the length of key
    /// @return child pointer
    inline void * Lookup(const char * piece, uint16_t piece_len) {
        assert(piece_len == trie_span);
        return mappings.Lookup(piece);
    }

    /// @brief Insert a new fanout into this trie node
    /// @param piece the insert key prefix, which is of type char *
    /// @param piece_len the size of the key, should be equal to trie_span
    /// @param child the insert child pointer
    /// @return whether the node should be rebuilt
    bool Insert(const char * piece, uint16_t piece_len, void * child) { 
        assert(piece_len == trie_span);
        keys.Insert(piece);
        if(!mappings.Insert(piece, child)) { 
            // rebuild the mappings structure due to size limit
            HashTable new_mappings(trie_span, mappings.capacity_ * RESIZE_FACTOR);
            for(int i = 0; i < count; i++) {
                void * child = mappings.Lookup(keys[i]);
                new_mappings.Insert(keys[i], child);
            }
            std::swap(mappings, new_mappings); // change to a new hashtable
        }
        count += 1;
        
        if(count >= MAX_FANOUT) {
            return false;
        }
    }
};

class LeafNode {
private:
    uint16_t type;
    uint16_t count;
    uint32_t used;
    uint32_t capacity;

    char * buffer;
    LeafNode * sibling;
    uint16_t slotarray[MAX_LEAF];

public:

    LeafNode(std::vector<char *> & keys, std::vector<void *> & vals, uint32_t start, uint32_t end) {
        count = 0;
        capacity = (end - start) * AVERAGE_KEYLEN;
        used = 0;
        buffer = new char[capacity];
        sibling = nullptr;

        for(int i = 0; i < keys.size(); i++) {
            Insert(keys[i], vals[start + i]);
        }
    }

    ~LeafNode() {
        delete [] buffer;
    }

    void * Lookup(const char * key) {
        return nullptr;
    }

    void Insert(const char * key, void * child) {
        if(count < MAX_LEAF) {
            int key_len = strlen(key);
            int total_len = key_len + 1 + sizeof(void *);
            if(used + total_len > capacity) {
                // enlarge the buffer
                int new_capacity = capacity + AVERAGE_KEYLEN * 8;
                char * new_buffer = new char[new_capacity];
                memcpy(new_buffer, buffer, used);
                std::swap(new_buffer, buffer);
                capacity = new_capacity;
                delete [] new_buffer;
            }

            // copy value into LeafNode
            memcpy(buffer + used, key, key_len);
            buffer[used + key_len] = 0;
            memcpy(buffer + key_len + 1, (&child), sizeof(void *));
            
            slotarray[count++] = used;
            used += total_len;
            count += 1;
        }
    }
}; 

int SampleDistinctPrefix(std::vector<char *> & keys, int split_span) {
    static const int SAMPLE_STEP = 32;
    // insert all possible prefixes into a set
    std::unordered_set<MyKey, MyKey::KeyHash, MyKey::KeyEqn> unique_keys;
    for(int i = 0; i < keys.size(); i += SAMPLE_STEP){
        unique_keys.emplace(split_span, keys[i]);
    }
    
    return unique_keys.size();
}

int GetSuboptimalSpan(std::vector<char *> & keys) {
    return 16;
}

void * Bulkload(std::vector<char *> & keys, std::vector<void *> & vals, uint32_t start, uint32_t end) {
    if(end - start + 1 < MAX_LEAF) {
        return new LeafNode(keys, vals, start, end);
    }
    
    int split_span = GetSuboptimalSpan(keys); 
    TrieNode * newnode = new TrieNode(split_span);

    // group all keys into several group, each group shares a common prefix of length split_span
    MyKey::KeyEqn eqn_func;
    for(uint32_t i = 0; i < keys.size(); ) {
        MyKey group_first_prefix = {split_span, keys[i]};
        std::vector<char *> group_keys{keys[i] + split_span};
        uint32_t group_start = start + i, group_end = start + i;

        // retrieve subkeys of each group
        while(++i < keys.size() && eqn_func(group_first_prefix, MyKey{split_span, keys[i]})) {
            group_keys.push_back(keys[i] + split_span);
            group_end += 1;
        }

        // build a new trie node for a group
        void * child = Bulkload(group_keys, vals, group_end, group_end);
        
        // add the child to root node
        newnode->Insert(group_first_prefix.str, split_span, child);
    }

    return newnode;
}